#include <X11/Xlib.h>
#include <X11/Xft/Xft.h>
#include <X11/keysym.h>
#include <X11/Xatom.h>
#include <security/pam_appl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <unistd.h>
#include <sys/mman.h>
#include <time.h>
#include <math.h>

// ================================================================================
// Configuration (adjust these as needed)
// ================================================================================
#define FONT_SIZE 28
#define HINT_FONT_SIZE 20
#define FONT_NAME "DejaVu Sans Mono"
#define MAX_PASSWORD_LEN 512
#define MAX_FAILED_ATTEMPTS 5
#define TARGET_FPS 60
#define FRAME_TIME_US (1000000 / TARGET_FPS)

// Colors (hex RGB)
#define COLOR_BG 0x000000        // Black background (fullscreen)
#define COLOR_INPUT_BG 0x282828  // Dark gray input box
#define COLOR_TEXT 0xFFFFFF      // White text
#define COLOR_PROMPT 0xFFCC00    // Amber prompt text
#define COLOR_ERROR 0xFF0000     // Red error text
#define COLOR_SUCCESS 0x00FF00   // Green success text
#define COLOR_SEPARATOR 0x505050 // Gray separator
#define COLOR_HINT 0x8A8A8A      // Light gray hints

// ================================================================================
// Global State
// ================================================================================
char password[MAX_PASSWORD_LEN] = {0};
size_t cursor = 0;
int input_len = 0;
int failed_attempts = 0;
bool is_locked = true;
bool show_help = false;
bool redraw_needed = true;

// X11 Resources
Display *dpy;
Window win;
GC gc;
Visual *vis;
Colormap cmap;
Pixmap backbuffer;
XftFont *main_font;
XftFont *hint_font;
int screen;
int screen_width, screen_height;

// Xft Colors
XftColor color_bg;
XftColor color_input_bg;
XftColor color_text;
XftColor color_prompt;
XftColor color_error;
XftColor color_success;
XftColor color_separator;
XftColor color_hint;

// ================================================================================
// PAM Authentication (unchanged)
// ================================================================================
static int pam_conv_callback(int num_msg, const struct pam_message **msg,
                             struct pam_response **resp, void *appdata_ptr)
{
    (void)appdata_ptr;
    *resp = calloc(num_msg, sizeof(struct pam_response));
    if (!*resp)
        return PAM_BUF_ERR;

    for (int i = 0; i < num_msg; i++)
    {
        if (msg[i]->msg_style == PAM_PROMPT_ECHO_OFF)
        {
            resp[i]->resp = strdup(password);
            resp[i]->resp_retcode = 0;
        }
    }
    return PAM_SUCCESS;
}

static bool verify_password()
{
    if (input_len == 0)
        return false;

    struct pam_conv conv = {pam_conv_callback, NULL};
    pam_handle_t *pam_handle = NULL;
    const char *username = getlogin();
    int ret = pam_start("login", username, &conv, &pam_handle);

    if (ret != PAM_SUCCESS)
    {
        fprintf(stderr, "PAM init failed: %s\n", pam_strerror(pam_handle, ret));
        pam_end(pam_handle, ret);
        return false;
    }

    ret = pam_authenticate(pam_handle, 0);
    pam_end(pam_handle, ret);

    if (ret == PAM_SUCCESS)
        return true;
    failed_attempts++;
    return false;
}

static void clear_password()
{
    memset(password, 0, sizeof(password));
    input_len = 0;
    cursor = 0;
    redraw_needed = true;
}

// ================================================================================
// X11 Initialization & Resources (MODIFIED FOR FULLSCREEN)
// ================================================================================
static bool init_colors()
{
    if (!XftColorAllocName(dpy, vis, cmap, "#000000", &color_bg))
        return false;
    if (!XftColorAllocName(dpy, vis, cmap, "#282828", &color_input_bg))
        return false;
    if (!XftColorAllocName(dpy, vis, cmap, "#FFFFFF", &color_text))
        return false;
    if (!XftColorAllocName(dpy, vis, cmap, "#FFCC00", &color_prompt))
        return false;
    if (!XftColorAllocName(dpy, vis, cmap, "#FF0000", &color_error))
        return false;
    if (!XftColorAllocName(dpy, vis, cmap, "#00FF00", &color_success))
        return false;
    if (!XftColorAllocName(dpy, vis, cmap, "#505050", &color_separator))
        return false;
    if (!XftColorAllocName(dpy, vis, cmap, "#8A8A8A", &color_hint))
        return false;
    return true;
}

static bool init_fonts()
{
    char main_font_str[256];
    snprintf(main_font_str, sizeof(main_font_str), "%s:pixelsize=%d", FONT_NAME, FONT_SIZE);
    main_font = XftFontOpenName(dpy, screen, main_font_str);
    if (!main_font)
    {
        fprintf(stderr, "Failed to load main font: %s\n", main_font_str);
        return false;
    }

    char hint_font_str[256];
    snprintf(hint_font_str, sizeof(hint_font_str), "%s:pixelsize=%d", FONT_NAME, HINT_FONT_SIZE);
    hint_font = XftFontOpenName(dpy, screen, hint_font_str);
    if (!hint_font)
    {
        fprintf(stderr, "Failed to load hint font: %s\n", hint_font_str);
        XftFontClose(dpy, main_font);
        return false;
    }
    return true;
}

// Helper to grab keyboard/mouse exclusively (blocks i3wm shortcuts)
static bool grab_input()
{
    // Grab keyboard (blocks all key events from reaching other windows)
    if (XGrabKeyboard(dpy, win, False, GrabModeAsync, GrabModeAsync, CurrentTime) != GrabSuccess)
    {
        fprintf(stderr, "Failed to grab keyboard (retrying...)\n");
        usleep(100000); // Retry after 100ms
        if (XGrabKeyboard(dpy, win, False, GrabModeAsync, GrabModeAsync, CurrentTime) != GrabSuccess)
        {
            fprintf(stderr, "Fatal: Could not grab keyboard\n");
            return false;
        }
    }

    // Grab mouse (blocks all mouse events)
    if (XGrabPointer(dpy, win, False, ButtonPressMask | ButtonReleaseMask | PointerMotionMask,
                     GrabModeAsync, GrabModeAsync, None, None, CurrentTime) != GrabSuccess)
    {
        fprintf(stderr, "Failed to grab pointer (retrying...)\n");
        usleep(100000);
        if (XGrabPointer(dpy, win, False, ButtonPressMask | ButtonReleaseMask | PointerMotionMask,
                         GrabModeAsync, GrabModeAsync, None, None, CurrentTime) != GrabSuccess)
        {
            fprintf(stderr, "Fatal: Could not grab pointer\n");
            XUngrabKeyboard(dpy, CurrentTime); // Cleanup
            return false;
        }
    }

    return true;
}

static bool init_x11()
{
    dpy = XOpenDisplay(NULL);
    if (!dpy)
    {
        fprintf(stderr, "Failed to open X display\n");
        return false;
    }

    screen = DefaultScreen(dpy);
    vis = DefaultVisual(dpy, screen);
    cmap = DefaultColormap(dpy, screen);
    screen_width = DisplayWidth(dpy, screen);   // Full screen width
    screen_height = DisplayHeight(dpy, screen); // Full screen height

    XSetWindowAttributes attr;
    attr.override_redirect = True; // Bypass window manager entirely
    attr.background_pixel = color_bg.pixel;
    attr.event_mask = ExposureMask | KeyPressMask | ButtonPressMask | PointerMotionMask;

    // Create FULLSCREEN window (covers entire screen)
    win = XCreateWindow(
        dpy, DefaultRootWindow(dpy),
        0, 0,                        // Top-left corner (covers entire screen)
        screen_width, screen_height, // Full screen dimensions
        0, DefaultDepth(dpy, screen),
        InputOutput, vis,
        CWOverrideRedirect | CWBackPixel | CWEventMask, &attr);

    // Set as fullscreen window (hints to window managers)
    Atom net_wm_state = XInternAtom(dpy, "_NET_WM_STATE", False);
    Atom net_wm_state_fullscreen = XInternAtom(dpy, "_NET_WM_STATE_FULLSCREEN", False);
    XChangeProperty(
        dpy, win, net_wm_state, XA_ATOM, 32,
        PropModeReplace, (unsigned char *)&net_wm_state_fullscreen, 1);

    // Create backbuffer for smooth rendering
    backbuffer = XCreatePixmap(dpy, win, screen_width, screen_height, DefaultDepth(dpy, screen));
    gc = XCreateGC(dpy, backbuffer, 0, NULL);
    if (!gc)
    {
        fprintf(stderr, "Failed to create GC\n");
        return false;
    }

    if (!init_colors() || !init_fonts())
        return false;

    XMapWindow(dpy, win);
    XSetInputFocus(dpy, win, RevertToParent, CurrentTime);
    XFlush(dpy);

    // Grab exclusive input (blocks i3wm shortcuts)
    if (!grab_input())
        return false;

    return true;
}

// ================================================================================
// Rendering Helpers (MODIFIED FOR FULLSCREEN LAYOUT)
// ================================================================================
static int text_width(const char *str, XftFont *font)
{
    if (!str || !*str)
        return 0;
    XGlyphInfo info;
    XftTextExtentsUtf8(dpy, font, (const FcChar8 *)str, strlen(str), &info);
    return info.width;
}

static void draw_rect(int x, int y, int w, int h, XftColor *color)
{
    XSetForeground(dpy, gc, color->pixel);
    XFillRectangle(dpy, backbuffer, gc, x, y, w, h);
}

static void draw_text(int x, int y, const char *str, XftColor *color, XftFont *font)
{
    if (!str || !*str)
        return;
    y += font->ascent; // Adjust for font ascent
    XftDraw *draw = XftDrawCreate(dpy, backbuffer, vis, cmap);
    XftDrawStringUtf8(draw, color, font, x, y, (const FcChar8 *)str, strlen(str));
    XftDrawDestroy(draw);
}

static void draw_cursor(int x, int y)
{
    y += main_font->ascent;
    draw_rect(x, y - main_font->height, 2, main_font->height, &color_text);
}

// ================================================================================
// Event Handling (unchanged)
// ================================================================================
static void handle_key_press(XKeyEvent *ev)
{
    if (!is_locked)
        return;

    KeySym keysym = XKeycodeToKeysym(dpy, ev->keycode, 0);
    char buf[32];
    int len = XLookupString(ev, buf, sizeof(buf), &keysym, NULL);
    bool ctrl_pressed = (ev->state & ControlMask) != 0;
    bool mod_pressed = (ev->state & Mod4Mask) != 0; // Block i3's mod key (usually Mod4)

    // Block all mod key combinations (critical for i3wm)
    if (mod_pressed)
        return;

    // Ctrl shortcuts
    if (ctrl_pressed)
    {
        switch (keysym)
        {
        case XK_a:
            cursor = 0;
            redraw_needed = true;
            return;
        case XK_e:
            cursor = input_len;
            redraw_needed = true;
            return;
        case XK_f:
            if (cursor < input_len)
            {
                cursor++;
                redraw_needed = true;
            }
            return;
        case XK_b:
            if (cursor > 0)
            {
                cursor--;
                redraw_needed = true;
            }
            return;
        case XK_d:
            if (cursor < input_len)
            {
                memmove(&password[cursor], &password[cursor + 1], MAX_PASSWORD_LEN - cursor - 1);
                input_len--;
                redraw_needed = true;
            }
            return;
        case XK_k:
            password[cursor] = '\0';
            input_len = cursor;
            redraw_needed = true;
            return;
        }
    }

    // Regular keys
    switch (keysym)
    {
    case XK_Escape:
        clear_password();
        return;
    case XK_Return:
    case XK_KP_Enter:
        if (verify_password())
            is_locked = false;
        else
            clear_password();
        redraw_needed = true;
        return;
    case XK_BackSpace:
        if (cursor > 0)
        {
            cursor--;
            memmove(&password[cursor], &password[cursor + 1], MAX_PASSWORD_LEN - cursor - 1);
            input_len--;
            redraw_needed = true;
        }
        return;
    case XK_Delete:
        if (cursor < input_len)
        {
            memmove(&password[cursor], &password[cursor + 1], MAX_PASSWORD_LEN - cursor - 1);
            input_len--;
            redraw_needed = true;
        }
        return;
    case XK_Right:
        if (cursor < input_len)
        {
            cursor++;
            redraw_needed = true;
        }
        return;
    case XK_Left:
        if (cursor > 0)
        {
            cursor--;
            redraw_needed = true;
        }
        return;
    case XK_Home:
        cursor = 0;
        redraw_needed = true;
        return;
    case XK_End:
        cursor = input_len;
        redraw_needed = true;
        return;
    case XK_F1:
        show_help = !show_help;
        redraw_needed = true;
        return;
    default:
        if (len > 0 && buf[0] >= 32 && buf[0] <= 126 && input_len < MAX_PASSWORD_LEN - 1)
        {
            memmove(&password[cursor + 1], &password[cursor], MAX_PASSWORD_LEN - cursor - 1);
            password[cursor++] = buf[0];
            input_len++;
            redraw_needed = true;
        }
    }
}

static void handle_events()
{
    XEvent ev;
    while (XPending(dpy) > 0 && is_locked)
    {
        XNextEvent(dpy, &ev);
        switch (ev.type)
        {
        case Expose:
            redraw_needed = true;
            break;
        case KeyPress:
            handle_key_press(&ev.xkey);
            break;
        case ButtonPress:
        case MotionNotify:
            redraw_needed = true;
            break; // Refresh on mouse movement
        }
    }
}

// ================================================================================
// Main Rendering (CENTERED UI IN FULLSCREEN)
// ================================================================================
static void render()
{
    if (!redraw_needed)
        return;

    // Clear fullscreen background
    draw_rect(0, 0, screen_width, screen_height, &color_bg);

    // Calculate centered UI dimensions
    const int popup_width = 800; // Keep input UI at a reasonable size
    const int popup_height = 300;
    const int popup_x = (screen_width - popup_width) / 2;   // Center horizontally
    const int popup_y = (screen_height - popup_height) / 2; // Center vertically

    // Draw semi-transparent overlay for the UI (optional but nicer)
    draw_rect(popup_x, popup_y, popup_width, popup_height, &color_bg);

    // Title
    const char *title = "Screen Locked";
    int title_x = popup_x + (popup_width - text_width(title, main_font)) / 2;
    draw_text(title_x, popup_y + 20, title, &color_text, main_font);

    // Input box
    int input_y = popup_y + 20 + FONT_SIZE + 20;
    draw_rect(popup_x + 20, input_y, popup_width - 40, FONT_SIZE + 12, &color_input_bg);

    // Prompt ("Password: ")
    const char *prompt = "Password: ";
    int prompt_x = popup_x + 28;
    draw_text(prompt_x, input_y + 6, prompt, &color_prompt, main_font);
    int prompt_w = text_width(prompt, main_font);

    // Password indicator (asterisks)
    char pass_indicator[MAX_PASSWORD_LEN + 1] = {0};
    memset(pass_indicator, '*', input_len);
    draw_text(prompt_x + prompt_w, input_y + 6, pass_indicator, &color_success, main_font);

    // Blinking cursor
    static double last_cursor_flip = 0;
    static bool cursor_visible = true;
    double now = (double)clock() / CLOCKS_PER_SEC;
    if (now - last_cursor_flip > 0.5)
    {
        cursor_visible = !cursor_visible;
        last_cursor_flip = now;
        redraw_needed = true;
    }
    if (cursor_visible)
    {
        char cursor_substr[MAX_PASSWORD_LEN];
        strncpy(cursor_substr, pass_indicator, cursor);
        cursor_substr[cursor] = '\0';
        int cursor_x = prompt_x + prompt_w + text_width(cursor_substr, main_font);
        draw_cursor(cursor_x, input_y + 6);
    }

    // Separator
    int separator_y = input_y + FONT_SIZE + 20;
    draw_rect(popup_x + 20, separator_y, popup_width - 40, 1, &color_separator);

    // Failed attempts
    if (failed_attempts > 0)
    {
        char fail_text[64];
        snprintf(fail_text, sizeof(fail_text), "Failed: %d/%d", failed_attempts, MAX_FAILED_ATTEMPTS);
        int fail_x = popup_x + (popup_width - text_width(fail_text, main_font)) / 2;
        draw_text(fail_x, separator_y + 15, fail_text, &color_error, main_font);

        if (failed_attempts >= MAX_FAILED_ATTEMPTS)
        {
            const char *block = "Temporarily Blocked";
            int block_x = popup_x + (popup_width - text_width(block, main_font)) / 2;
            draw_text(block_x, separator_y + 15 + FONT_SIZE, block, &color_error, main_font);
        }
    }

    // Help text (F1)
    if (show_help)
    {
        const char *help[] = {
            "ESC: Clear | Enter: Submit | F1: Hide Help",
            "Ctrl+A: Home | Ctrl+E: End | Ctrl+F/B: Left/Right",
            "Backspace: Delete | Ctrl+D: Delete Forward | Ctrl+K: Clear Rest"};
        int help_y = separator_y + 40;
        for (int i = 0; i < 3; i++)
        {
            int help_x = popup_x + (popup_width - text_width(help[i], hint_font)) / 2;
            draw_text(help_x, help_y + i * (HINT_FONT_SIZE + 10), help[i], &color_hint, hint_font);
        }
    }
    else
    {
        const char *hint = "Press F1 for help";
        int hint_x = popup_x + popup_width - 20 - text_width(hint, hint_font) - 5;
        draw_text(hint_x, popup_y + popup_height - 20 - HINT_FONT_SIZE, hint, &color_hint, hint_font);
    }

    // Swap buffers
    XCopyArea(dpy, backbuffer, win, gc, 0, 0, screen_width, screen_height, 0, 0);
    XFlush(dpy);
    redraw_needed = false;
}

// ================================================================================
// Cleanup (RELEASE GRABS)
// ================================================================================
static void cleanup()
{
    // Release input grabs
    XUngrabKeyboard(dpy, CurrentTime);
    XUngrabPointer(dpy, CurrentTime);

    // Securely clear password
    memset(password, 0, sizeof(password));

#ifdef __linux__
    munlock(password, sizeof(password));
#endif

    // Free X11 resources
    XftFontClose(dpy, main_font);
    XftFontClose(dpy, hint_font);

    XftColorFree(dpy, vis, cmap, &color_bg);
    XftColorFree(dpy, vis, cmap, &color_input_bg);
    XftColorFree(dpy, vis, cmap, &color_text);
    XftColorFree(dpy, vis, cmap, &color_prompt);
    XftColorFree(dpy, vis, cmap, &color_error);
    XftColorFree(dpy, vis, cmap, &color_success);
    XftColorFree(dpy, vis, cmap, &color_separator);
    XftColorFree(dpy, vis, cmap, &color_hint);

    XFreePixmap(dpy, backbuffer);
    XFreeGC(dpy, gc);
    XDestroyWindow(dpy, win);
    XCloseDisplay(dpy);
}

// ================================================================================
// Main
// ================================================================================
int main()
{
#ifdef __linux__
    if (mlock(password, sizeof(password)) != 0)
    {
        fprintf(stderr, "Warning: Could not lock password memory (run with sufficient privileges)\n");
    }
#endif

    if (!init_x11())
    {
        cleanup();
        return 1;
    }

    // Main loop
    struct timespec last_frame;
    clock_gettime(CLOCK_MONOTONIC, &last_frame);
    while (is_locked)
    {
        handle_events();
        render();

        // Throttle FPS
        struct timespec now;
        clock_gettime(CLOCK_MONOTONIC, &now);
        long elapsed = (now.tv_sec - last_frame.tv_sec) * 1000000 +
                       (now.tv_nsec - last_frame.tv_nsec) / 1000;
        if (elapsed < FRAME_TIME_US)
        {
            usleep(FRAME_TIME_US - elapsed);
        }
        last_frame = now;
    }

    cleanup();
    return 0;
}
